import sys
sys.path.append(".")

import os
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

from dataset_tools.config import get_config 
from dataset_tools.gen_random_mjcf_scenes import GenMJCFScenes
from dataset_tools.trans_xml2pngs import MJCF2GRID
from dataset_tools.gen_dataset import GenDataset
from dataset_tools.collect_simple_data import RunMujoco, load_csv_to_dict_list
from dataset_tools.sample_data_post_process import process_images


config = get_config()


print("Generating scenes...")

stairwall_cfg = config["Gen_scene"].get("stairwall", {})
if stairwall_cfg:
    genStairwallMJCFScenes = GenMJCFScenes(
        mode="stairwell",
        collision_num=stairwall_cfg["collision_num"],
        gen_scene_num=stairwall_cfg["gen_scene_num"],
        scene_save_dir=config["Gen_scene"]["scene_save_dir"],
    )
    genStairwallMJCFScenes.gen()

random_cfg = config["Gen_scene"].get("random", {})
if random_cfg:
    genRandomMJCFScenes = GenMJCFScenes(
        mode="random",
        collision_num=random_cfg["collision_num"],
        gen_scene_num=random_cfg["gen_scene_num"],
        scene_save_dir=config["Gen_scene"]["scene_save_dir"],
    )
    genRandomMJCFScenes.gen()



print("Generating grid images...")
scenes_dir_list = [
    dir 
    for dir in os.listdir(config["Gen_scene"]["scene_save_dir"]) 
    if os.path.isdir(os.path.join(config["Gen_scene"]["scene_save_dir"], dir)) and "scene_" in dir
]
for scene_dir in scenes_dir_list:
    mjcf2grid = MJCF2GRID(os.path.join(config["Gen_scene"]["scene_save_dir"], scene_dir), resolution=0.02)
    mjcf2grid.run()


print("Generating dataset [real]...")
genDataset = GenDataset(
    scenes_dir = config["Gen_scene"]["scene_save_dir"],
    dataset_save_dir = config["Gen_dataset"]["dataset_save_dir"],
    resolution =config["Gen_dataset"]["resolution"],
    dataset_pre_scene = config["Gen_dataset"]["dataset_pre_scene"]
)
genDataset.run()

print("Generating dataset [sample]...")
image_data_csv = os.path.join(config["Gen_dataset"]["dataset_save_dir"], "image_data.csv")
image_data = load_csv_to_dict_list(image_data_csv)
sample_dir = os.path.join(config["Gen_dataset"]["dataset_save_dir"], "sample")

def process(data):
    scene_dir = os.path.join(config["Gen_scene"]["scene_save_dir"], data["Scene"])
    runMujoco = RunMujoco(
        data,
        scene_dir,
        config["Collect_sim_data"]["laser_path"], 
        output_dir = sample_dir
    )
    return runMujoco

with ThreadPoolExecutor(max_workers=config["Collect_sim_data"]["worker"]) as executor:
    results = list(tqdm(executor.map(process, image_data), total=len(image_data), desc="Processing"))

print("Post processing sample data...")
real_dir = os.path.join(config["Gen_dataset"]["dataset_save_dir"], "real")
process_images(
    input_dir=real_dir,
    output_dir=sample_dir,
    speckle_prob=config["Sample_data_post_process"]["speckle_prob"],
    speckle_size_lower=config["Sample_data_post_process"]["speckle_size_lower"],
    speckle_size_upper=config["Sample_data_post_process"]["speckle_size_upper"]
)